from qdrant_client import QdrantClient, models
from sentence_transformers import SentenceTransformer
from typing import List, Dict, Optional
import numpy as np
import os

os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'


class BGE_M3_QdrantIndexer:
    def __init__(self, collection_name: str = "zy_collection", host: str = "192.168.77.101", port: int = 6333):
        """
        初始化 BGE-M3 和 Qdrant 客户端

        参数:
            collection_name: Qdrant 集合名称
            host: Qdrant 服务器地址
            port: Qdrant 服务器端口
        """
        # 初始化 BGE-M3 模型
        self.model = SentenceTransformer(r'I:\models\BAAI\bge-m3')


        # self.client = QdrantClient(
        #     url="https://192.168.77.101",
        #     api_key="123456",
        #     https=True,
        #     ssl_context=ssl_context
        # )

        self.client = QdrantClient(
            url="http://192.168.77.101:6333",
            api_key="123456"
        )

        # 初始化 Qdrant 客户端
        # self.client = QdrantClient(host=host, port=port,api_key='123456',
        #                            verify=False )

        # 集合配置
        self.collection_name = collection_name
        self.vector_size = self.model.get_sentence_embedding_dimension()

        # 检查并创建集合
        self._ensure_collection_exists()

    def _ensure_collection_exists(self):
        """确保集合存在，如果不存在则创建"""
        try:
            collection_info = self.client.get_collection(self.collection_name)
            print(f"使用现有集合: {self.collection_name}")
        except Exception:
            print(f"创建新集合: {self.collection_name}")
            self.client.create_collection(
                collection_name=self.collection_name,
                vectors_config=models.VectorParams(
                    size=self.vector_size,
                    distance=models.Distance.COSINE
                )
            )

    def encode_text(self, text: str) -> Dict:
        """
        使用 BGE-M3 编码文本，返回密集向量、稀疏向量和 ColBERT 令牌

        参数:
            text: 要编码的文本

        返回:
            包含 'dense'、'sparse' 和 'colbert' 向量的字典
        """
        # 获取所有类型的嵌入
        output = self.model.encode(text, return_dense=True, return_sparse=True, return_colbert_vecs=True)

        return output
        # return {
        #     "dense": output['dense_vecs'].flatten().tolist(),
        #     "sparse": {
        #         "indices": output['sparse_vecs']['indices'].flatten().tolist(),
        #         "values": output['sparse_vecs']['values'].flatten().tolist()
        #     },
        #     "colbert": output['colbert_vecs'].tolist()
        # }

    def index_documents(self, documents: List[Dict], batch_size: int = 100):
        """
        索引文档到 Qdrant

        参数:
            documents: 文档列表，每个文档应包含 "id" 和 "text" 字段
            batch_size: 批量处理大小
        """
        points = []

        for doc in documents:
            # 编码文本
            vectors = self.encode_text(doc["text"])

            # 构建 Qdrant 点结构
            point = models.PointStruct(
                id=doc["id"],
                vector= vectors,
                # {
                #     "dense": vectors["dense"],
                #     "sparse": vectors["sparse"]
                # },
                payload={
                    "text": doc["text"],
                    "metadata": doc.get("metadata", {})
                }
            )
            points.append(point)

            # 批量提交
            if len(points) >= batch_size:
                self.client.upsert(
                    collection_name=self.collection_name,
                    points=points
                )
                points = []
                print(f"已索引 {len(points)} 文档")

        # 提交剩余文档
        if points:
            self.client.upsert(
                collection_name=self.collection_name,
                points=points
            )
            print(f"已索引 {len(points)} 文档 (最后一批)")

    def search(self, query: str, top_k: int = 5, use_hybrid: bool = True) -> List[Dict]:
        """
        搜索相似文档

        参数:
            query: 查询文本
            top_k: 返回的结果数量
            use_hybrid: 是否使用混合搜索

        返回:
            搜索结果列表
        """
        # 编码查询
        query_vectors = self.encode_text(query)

        # 构建搜索请求
        if use_hybrid:

            # 混合搜索 (密集 + 稀疏)
            search_result = self.client.search(
                collection_name=self.collection_name,
                query_vector=query_vectors,
                # ,
                # query_sparse_vector=models.NamedSparseVector(
                #     name="sparse",
                #     vector=models.SparseVector(
                #         indices=query_vectors["sparse"]["indices"],
                #         values=query_vectors["sparse"]["values"]
                #     )
                # ),
                limit=top_k,
                with_payload=True,
                with_vectors=False
            )
        else:
            # 仅密集向量搜索
            search_result = self.client.search(
                collection_name=self.collection_name,
                query_vector=query_vectors,
                # query_vector=models.NamedVector(
                #     name="dense",
                #     vector=query_vectors["dense"]
                # ),
                limit=top_k,
                with_payload=True,
                with_vectors=False
            )

        # 格式化结果
        results = []
        for hit in search_result:
            results.append({
                "id": hit.id,
                "score": hit.score,
                "text": hit.payload,
                "metadata": hit.payload.get("metadata", {})
            })

        return results

    def delete_collection(self):
        """删除当前集合"""
        self.client.delete_collection(self.collection_name)
        print(f"集合 {self.collection_name} 已删除")


# 示例使用
if __name__ == "__main__":
    # 初始化索引器
    indexer = BGE_M3_QdrantIndexer()

    # # 示例文档
    documents = [
        {"id": 1, "text": "深度学习是机器学习的一个子领域", "metadata": {"category": "AI"}},
        {"id": 2, "text": "Python是一种流行的编程语言", "metadata": {"category": "编程"}},
        {"id": 3, "text": "Qdrant是一个高效的向量搜索引擎", "metadata": {"category": "数据库"}},
        {"id": 4, "text": "BGE-M3是BAAI开发的多功能嵌入模型", "metadata": {"category": "NLP"}},
        {"id": 5, "text": "大语言模型如GPT-4正在改变AI领域", "metadata": {"category": "AI"}},
    ]

    # 索引文档
    print("开始索引文档...")
    indexer.index_documents(documents)
    print("索引完成!")

    # 执行搜索
    query = "人工智能技术"
    print(f"\n搜索查询: '{query}'")
    results = indexer.search(query)

    print("\n搜索结果:")
    for i, result in enumerate(results, 1):
        print(f"{i}. ID: {result['id']}, 分数: {result['score']:.4f}")
        print(f"   文本: {result['text']}")
        print(f"   元数据: {result['metadata']}\n")

    # 清理 (可选)
    # indexer.delete_collection()